import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# 设置随机种子以确保可重复性
torch.manual_seed(42)

# 定义 LeNet-5 模型
class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5, stride=1, padding=2,)  # 卷积层1
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)  # 池化层1
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5, stride=1)  # 卷积层2
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)  # 池化层2
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 全连接层1
        self.fc2 = nn.Linear(120, 84)  # 全连接层2
        self.fc3 = nn.Linear(84, 10)  # 输出层

    def forward(self, x):
        x = self.pool1(torch.relu(self.conv1(x)))  # 卷积 -> ReLU -> 池化
        x = self.pool2(torch.relu(self.conv2(x)))  # 卷积 -> ReLU -> 池化
        x = x.view(-1, 16 * 5 * 5)  # 展平
        x = torch.relu(self.fc1(x))  # 全连接 -> ReLU
        x = torch.relu(self.fc2(x))  # 全连接 -> ReLU
        x = self.fc3(x)  # 输出层
        return x

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),  # 转换为Tensor
    transforms.Normalize((0.1307,), (0.3081,))  # 归一化
])

# 下载 MNIST 数据集（使用清华镜像）
train_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=True,
    download=True,
    transform=transform
)
test_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=False,
    download=True,
    transform=transform
)

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 初始化模型、损失函数和优化器
model = LeNet5()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
def train(model, train_loader, criterion, optimizer, num_epochs=5):
    model.train()
    for epoch in range(num_epochs):
        running_loss = 0.0
        for i, (inputs, labels) in enumerate(train_loader):
            optimizer.zero_grad()  # 梯度清零
            outputs = model(inputs)  # 前向传播
            loss = criterion(outputs, labels)  # 计算损失
            loss.backward()  # 反向传播
            optimizer.step()  # 更新参数

            running_loss += loss.item()
            if (i + 1) % 100 == 0:
                print(f"Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{len(train_loader)}], Loss: {loss.item():.4f}")

# 测试模型
def test(model, test_loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in test_loader:
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print(f"Test Accuracy: {100 * correct / total:.2f}%")

# 训练和测试
train(model, train_loader, criterion, optimizer, num_epochs=5)
test(model, test_loader)